# -*- coding: utf-8 -*-
# time: 2025/4/17 09:29
# file: hf_Ollama_LangChain01.py
# author: hanson
"""
通过 Hugging Face 获取模型 → 用 Ollama 本地化部署 → 通过 LangChain 集成到业务系统中
"""
import os
from typing import Dict, Any
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
from langchain.chains import LLMChain
from huggingface_hub import hf_hub_download


# 1. 从 Hugging Face 获取模型
def download_model_from_hf(repo_id: str, model_filename: str, local_dir: str = "./models") -> str:
    """
    从 Hugging Face 下载模型文件

    参数:
        repo_id: Hugging Face 仓库ID (如 "TheBloke/Llama-2-7B-GGML")
        model_filename: 模型文件名 (如 "llama-2-7b.ggmlv3.q4_0.bin")
        local_dir: 本地保存目录

    返回:
        模型本地路径
    """
    os.makedirs(local_dir, exist_ok=True)
    model_path = hf_hub_download(
        repo_id=repo_id,
      #  filename=model_filename,
        local_dir=local_dir,
        local_dir_use_symlinks=False
    )
    print(f"模型已下载到: {model_path}")
    return model_path


# 2. 使用 Ollama 本地部署模型
def setup_ollama_model(model_path: str, model_name: str = "custom-llm"):
    """
    设置 Ollama 本地模型

    参数:
        model_path: 模型文件路径
        model_name: 在 Ollama 中使用的模型名称
    """
    import ollama

    # 检查模型是否已存在
    try:
        ollama.show(model_name)
        print(f"模型 '{model_name}' 已存在，跳过创建")
        return
    except ollama.ResponseError:
        pass

    # 创建 Ollama 模型
    modelfile = f"""
    FROM {model_path}
    """

    ollama.create(model=model_name, modelfile=modelfile)
    print(f"Ollama 模型 '{model_name}' 创建成功")


# 3. 通过 LangChain 集成
def create_qa_chain(model_name: str = "custom-llm") -> LLMChain:
    """
    创建问答链

    参数:
        model_name: Ollama 模型名称

    返回:
        LangChain LLMChain 对象
    """
    # 初始化 Ollama LLM
    llm = Ollama(
        model=model_name,
        temperature=0.7,
        top_p=0.9,
        repeat_penalty=1.1
    )

    # 创建提示模板
    prompt_template = """
    你是一个专业的客服助手，请根据以下上下文回答问题。
    保持回答简洁专业，如果不知道答案，就说你不知道。

    上下文: {context}
    问题: {question}

    回答:
    """

    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )

    # 创建链
    qa_chain = prompt | llm     # LLMChain(llm=llm, prompt=prompt)
    return qa_chain


# 4. 业务系统集成
class CustomerServiceSystem:
    def __init__(self, qa_chain: LLMChain):
        self.qa_chain = qa_chain
        self.context_db = {
            "退货政策": "我们提供30天无理由退货服务，商品需保持原状。",
            "配送时间": "标准配送需要3-5个工作日，加急配送1-2个工作日。",
            "支付方式": "我们接受信用卡、支付宝、微信支付和银行转账。"
        }

    def get_context(self, topic: str) -> str:
        """从数据库获取上下文"""
        return self.context_db.get(topic, "没有找到相关信息")

    def handle_query(self, question: str) -> str:
        """处理客户查询"""
        # 简单识别问题主题（实际应用中可以使用更复杂的NLP技术）
        topic = self.identify_topic(question)
        context = self.get_context(topic)

        # 使用 LangChain 处理问题
        # response = self.qa_chain.run(
        #     context=context,
        #     question=question
        # )
        # 使用 LCEL 方式调用
        response = self.qa_chain.invoke({
            "context": context,
            "question": question
        })

        return response

    def identify_topic(self, question: str) -> str:
        """简单识别问题主题"""
        question = question.lower()
        if "退货" in question or "退款" in question:
            return "退货政策"
        elif "配送" in question or "发货" in question or "运输" in question:
            return "配送时间"
        elif "支付" in question or "付款" in question:
            return "支付方式"
        return "general"


# 主程序
def main():
    # 配置参数
    HF_REPO_ID = r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct"
    HF_MODEL_FILE = "llama-2-7b.ggmlv3.q4_0.bin"
    OLLAMA_MODEL_NAME = "qwen2.5:0.5b"

    try:
        # 1. 从 Hugging Face 下载模型
        print("正在从 Hugging Face 下载模型...")
        # model_path = download_model_from_hf(
        #     repo_id=HF_REPO_ID,
        #     model_filename=HF_MODEL_FILE
        # )
        model_path = r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct"
        # 2. 设置 Ollama 模型
        print("\n设置 Ollama 模型...")
        setup_ollama_model(model_path, model_name=OLLAMA_MODEL_NAME)

        # 3. 创建 LangChain 问答链
        print("\n创建 LangChain 问答链...")
        qa_chain = create_qa_chain(model_name=OLLAMA_MODEL_NAME)

        # 4. 集成到业务系统
        print("\n初始化客服系统...")
        css = CustomerServiceSystem(qa_chain)

        # 测试查询
        queries = [
            "你们的退货政策是什么？",
            "订单一般多久能发货？",
            "可以用微信支付吗？",
            "你们有实体店吗？"
        ]

        print("\n测试客服系统:")
        for query in queries:
            print(f"\n客户问题: {query}")
            response = css.handle_query(query)
            print(f"助手回答: {response}")

    except Exception as e:
        print(f"发生错误: {str(e)}")


if __name__ == "__main__":
    main()